import PIL.Image as Image 
import numpy as np
from numpy.testing._private.utils import tempdir 
from gauseKernel import Conv2D
from EdgeDetect import edgeDetectKernels


class Conv2DWithGradDetection(Conv2D):
    def __init__(self, kernel_type, sigma = 1, boundary = "zero", padding = "SAME"):
        super().__init__(1, 3, sigma, boundary=boundary, padding=padding)
        xKernel, yKernel = edgeDetectKernels(kernel_type, sigma)
        self.filters = [xKernel, yKernel]
      

class CornerPointDetector():
    def __init__(self, image, name, alpha = 0.04, kernel_type = "sobel", sigma = 1, boundary = "copy", padding = "SAME"):
        self.conv = Conv2DWithGradDetection(kernel_type, sigma, boundary, padding)
        self.image = image
        self.alpha = alpha
        self.name = name
    def getGrads(self):
        inputs = np.array(self.image)
        afterFliter = self.conv(inputs)
        gauseFilter = Conv2D(1, 5, 1, "copy", "SAME")
        gauseFilter1 = Conv2D(1, 5, 1, "copy", "SAME")
        gauseFilter2 = Conv2D(1, 5, 1, "copy", "SAME")
        xGrad = afterFliter[0]
        yGrad = afterFliter[1]
        self.xGrad = xGrad
        self.yGrad = yGrad
        self.xxGrad = self.xGrad ** 2 / 255
        # xxGImage = Image.fromarray(self.xxGrad.astype("uint"))
        # xxGImage.show()
        self.yyGrad = self.yGrad ** 2 / 255
        # yyGImage = Image.fromarray(self.yyGrad.astype("uint"))
        # yyGImage.show()
        self.xyGrad = self.xGrad * self.yGrad / 255
        # xyGImage = Image.fromarray(self.xyGrad.astype("uint"))
        # xyGImage.show()
        self.gxxGrad = gauseFilter(self.xxGrad)[0]
        # gxxGImage = Image.fromarray(self.gxxGrad.astype("uint"))
        # gxxGImage.show()
        self.gyyGrad = gauseFilter1(self.yyGrad)[0]
        # gyyGImage = Image.fromarray(self.gyyGrad.astype("uint"))
        # gyyGImage.show()
        self.gxyGrad = gauseFilter2(self.xyGrad)[0]
        # gxyGImage = Image.fromarray(self.gxyGrad.astype("uint"))
        # gxyGImage.show()
    def getHarris(self):
        self.har = (self.gxxGrad * self.gyyGrad - self.gxyGrad ** 2 - self.alpha * (self.gxxGrad + self.gyyGrad) ** 2)
        harImage = Image.fromarray(self.har.astype("uint")).convert("L")
        harImage.save("./CornerPointDetect/cornerPoint-" + self.name + "-" + str(self.alpha) + ".jpg")
    def nonMaxSuppression(self):
        h = len(self.har)
        w = len(self.har[0])
        tmphar = np.array(self.har)
        xchange = [-1, 0, 1, 0]
        ychange = [0, 1, 0, -1]
        for i in range(h):
            for j in range(w):
                for k in range(4):
                    if i + xchange[k] >= 0 and i + xchange[k] < h and j + ychange[k] >= 0 and j + ychange[k] < w:
                        if self.har[i][j] < self.har[i + xchange[k]][j + ychange[k]]:
                            tmphar[i][j] = 0
                            break
        NMSImage = Image.fromarray(tmphar.astype("uint")).convert("L")
        NMSImage.save("./CornerPointDetect/cornerPoint-" + self.name + "-" + str(self.alpha) + "-NMS.jpg")
    def run(self):
        self.getGrads()
        self.getHarris()
        self.nonMaxSuppression()

image = Image.open("blackwhiteL.jpg")
cornerDetector = CornerPointDetector(image, "blackwhite", alpha=0.04)
cornerDetector.run()

image = Image.open("rawImage.jpg")
cornerDetector2 = CornerPointDetector(image, "cat", alpha=0.04)
cornerDetector2.run()
    


